#!/bin/bash

SETDIR="/tmp/dispatch_cpu/setdir"
MEMDIR="/tmp/dispatch_cpu/memdir"
SHAREDIR="/tmp/dispatch_cpu/sharedir"
TOOL_CPUSHARE_DEFAULT=8

declare -A cpu_set
cpu_set=(
[umb_interface.dag]="6-11"
[acu_manager]="6-11"
[topmanagement_master.dag]="6-11"
[gaia_diagnostic.dag]="6-11"
[gaia_net_service.dag]="6-11"
[anp-soa-master.dag]="6-11"
[patrol_service.dag]="6-11"
[DataClosedLoop_master.dag]="6-11"
[pslam]="6-11"
[config_server]="6-11"
[SoaMiningService_master.dag]="6-11"
[l2c_service_master.dag]="6-11"
#[canbus_proxy_jidu.anvp.dag]="6-11"
[patrol_orinx_master_sched]="6-11"
[dreamview_sched]="6-11"
[topmanagement_slave.dag]="6-11"
[gaia_net_service.dag]="6-11"
[anp-soa-slave.dag]="6-11"
[DataClosedLoop_slave.dag]="6-11"
[aeb_fusion_online.dag]="6-11"
[podom.dag]="6-11"
[repo_sentry_manager.dag]="6-11"
[SoaMiningService_slave.dag]="6-11"
[l2c_service_slave.dag]="6-11"
)

i=0
for pname in $(echo ${!cpu_set[*]});
do
  process_name[$i]=${pname}
  process_cpuset[$i]=${cpu_set[${pname}]}
  process_cpushare[$i]=${TOOL_CPUSHARE_DEFAULT}
  process_type[$i]="cpu"
  ((++i))
done

# declare -a tool_set=(
# "x265_decompress"
# )

# for pname in "${tool_set[@]}"
# do
#   process_name[$i]=${pname}
#   process_cpuset[$i]=${TOOL_CPUSET_DEFAULT}
#   process_cpushare[$i]=${TOOL_CPUSHARE_DEFAULT}
#   process_type[$i]="tool"
#   ((++i))
# done

# setcpuset taskgrp pids cpuset
setcpuset() {
  taskgrp=$1
  cpuset=$3

  # create task group
  mkdir -p ${SETDIR}/${taskgrp}
  # set mem node
  echo 0 > ${SETDIR}/${taskgrp}/cpuset.mems
  # attach tasks
  pids=$2
  for pid in `echo $pids | awk '{print $0}'`
  do
    echo ${pid} > ${SETDIR}/${taskgrp}/tasks
  done
  # cpuset
  echo $cpuset > ${SETDIR}/${taskgrp}/cpuset.cpus
}

# setcpushare taskgrp pids share
setcpushare() {
  taskgrp=$1
  share=$3

  # create task group
  mkdir -p ${SHAREDIR}/${taskgrp}
  # reset
  echo 0 > ${SHAREDIR}/${taskgrp}/tasks
  # attach tasks
  pids=$2
  for pid in `echo $pids | awk '{print $0}'`
  do
    echo ${pid} > ${SHAREDIR}/${taskgrp}/tasks
  done
  # set share val
  echo $3 > ${SHAREDIR}/${taskgrp}/cpu.shares
}

function dispatch_e2e() {
  PROCESS_IDS=`ps -eo pid,cmd | grep "$1" | grep -v grep | awk '{print $1}'`
  setcpuset $1 "$PROCESS_IDS" $2
#   setcpushare $1 "$PROCESS_IDS" $3
  for PROCESS_ID in `echo $PROCESS_IDS | awk '{print $0}'`
  do
    #taskset -pc ${PROCESS_ID}
    SUBPROCESS_IDS=`ps -eT | grep -w  ${PROCESS_ID} | awk '{print $2}'`
    setcpuset $1 "${SUBPROCESS_IDS}" $2
    # setcpushare $1 "${SUBPROCESS_IDS}" $3
    echo $1
    #for i in `ps -eT | grep -w  ${PROCESS_ID} | awk '{print $2}'`; do taskset -pc $i; done
  done
}

function dispatch_tool() {
  PROCESS_IDS=`ps -eo pid,cmd | grep "$1" | grep -v grep | awk '{print $1}'`
  setcpuset tool "$PROCESS_IDS" $2
  setcpushare tool "$PROCESS_IDS" $3
  for PROCESS_ID in `echo $PROCESS_IDS | awk '{print $0}'`
  do
    #taskset -pc ${PROCESS_ID}
    SUBPROCESS_IDS=`ps -eT | grep -w  ${PROCESS_ID} | awk '{print $2}'`
    setcpuset tool "${SUBPROCESS_IDS}" $2
    setcpushare tool "${SUBPROCESS_IDS}" $3
    echo $1
    #for i in `ps -eT | grep -w  ${PROCESS_ID} | awk '{print $2}'`; do taskset -pc $i; done
  done
}

function dispatch_cpu() {
  mkdir -p ${SETDIR} > /dev/null
  sudo mount -t cgroup -o cpuset cpuset ${SETDIR}
#   sudo mount -t cgroup -o memory memory ${SETDIR}
#   mkdir -p ${SHAREDIR} > /dev/null
#   sudo mount -t cgroup -o cpu,cpuacct cpu,cpuacct ${SHAREDIR}

  PROCESS_NUM=${#process_name[*]}
  idx=0
  while [[ $idx -lt ${PROCESS_NUM} ]]; do
    echo $idx
    NAME=${process_name[${idx}]}
    CPUSET=${process_cpuset[${idx}]}
    CPUSHARE=${process_cpushare[${idx}]}
    TYPE=${process_type[${idx}]}
    if [[ "${TYPE}" == "cpu" ]]; then
      echo "dispatch_e2e ${NAME} ${CPUSET} ${CPUSHARE}"
      dispatch_e2e ${NAME} ${CPUSET} ${CPUSHARE}
    fi
    if [[ "${TYPE}" == "tool" ]]; then
      echo "dispatch_tool ${NAME} ${CPUSET} ${CPUSHARE}"
      dispatch_tool ${NAME} ${CPUSET} ${CPUSHARE}
    fi
    ((idx++))
  done
  #show_cpu_dispatch
}

function show_one() {
  PROCESS_IDS=`ps -eo pid,cmd | grep "$1" | grep -v grep | awk '{print $1}'`
  for PROCESS_ID in `echo $PROCESS_IDS | awk '{print $0}'`
  do
    taskset -pc ${PROCESS_ID}
    SUBPROCESS_IDS=`ps -eT | grep -w  ${PROCESS_ID} | awk '{print $2}'`
    echo $1
    for i in `ps -eT | grep -w  ${PROCESS_ID} | awk '{print $2}'`; do taskset -pc $i; done
  done
}

function show_cpu_dispatch() {
  PROCESS_NUM=${#process_name[*]}
  idx=0
  while [[ $idx -lt ${PROCESS_NUM} ]]; do
    NAME=${process_name[${idx}]}
    echo "${idx} : ${NAME}"
    CPUSET=${process_cpuset[${idx}]}
    CPUSHARE=${process_cpushare[${idx}]}
    TYPE=${process_type[${idx}]}
    show_one ${NAME} ${CPUSET} ${CPUSHARE}
    ((idx++))
  done
  find ${SHAREDIR} -name cpu.shares | xargs awk '{print FILENAME,$0}'
  wc -l ${SHAREDIR}/*/tasks
}

function recovery_cpu() {
  echo "recover"
  sudo umount ${SETDIR}
}

if [ $# -eq 0 ]; then
  dispatch_cpu
else
  case $1 in
    dispatch)
      dispatch_cpu
      ;;
    recovery)
      recovery_cpu
      ;;
    show)
      show_cpu_dispatch
      ;;
    help)
      echo "Useage:"
      echo "       bash $0 help : show useage."
      echo "       sudo bash $0 : dispatch cpu for program"
      echo "       sudo bash $0 dispatch : dispatch cpu for program"
      echo "       sudo bash $0 recovery : recovery cpu dispatch"
      echo "       sudo bash $0 show: show cpu dispatch"
      ;;
    *)
      echo "ERR: Unknown parameter: $1"
      exit 1
      ;;
  esac
fi